// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <random>
#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>

namespace kumo::pollux::functions {
    namespace {
        // See documentation at
        // https://prestodb.io/docs/current/functions/array.html#shuffle
        //
        // This function will shuffle identical arrays independently, i.e. even when
        // the input has duplicate rows represented using constant and dictionary
        // encoding, the output is flat and likely yields different values.
        //
        // E.g.1: constant encoding
        // Input: ConstantVector(base=ArrayVector[{1,2,3}], length=3, index=0)
        // Possible Output: ArrayVector[{1,3,2},{2,3,1},{3,2,1}]
        //
        // E.g.2: dict encoding
        // Input: DictionaryVector(
        //   dictionaryValues=ArrayVector[{1,2,3},{4,5},{1,2,3}],
        //   dictionaryIndices=[1,2,0])
        // Possible Output: ArrayVector[{5,4},{2,1,3},{1,3,2}]
        //
        class ArrayShuffleFunction : public exec::VectorFunction {
        public:
            explicit ArrayShuffleFunction(int32_t seed)
                : randGen_(std::make_unique<std::mt19937>(seed)) {
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /*outputType*/,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                POLLUX_CHECK_GE(args.size(), 1);

                // This is a non-deterministic function, which violates the guarantee on a
                // deterministic single-arg function that the expression evaluation will
                // peel off encodings, and we will only see flat or constant inputs. Hence,
                // we need to use DecodedVector to handle ALL encodings.
                exec::DecodedArgs decodedArgs(rows, args, context);
                auto decodedArg = decodedArgs.at(0);
                auto array_vector = decodedArg->base()->as<ArrayVector>();
                auto elementsVector = array_vector->elements();

                vector_size_t numElements = 0;
                context.applyToSelectedNoThrow(rows, [&](auto row) {
                    const auto size = array_vector->sizeAt(decodedArg->index(row));
                    numElements += size;
                });

                // Allocate new buffer to hold shuffled indices.
                BufferPtr shuffledIndices = allocate_indices(numElements, context.pool());
                BufferPtr offsets = allocateOffsets(rows.end(), context.pool());
                BufferPtr sizes = allocateSizes(rows.end(), context.pool());

                vector_size_t *rawIndices = shuffledIndices->asMutable<vector_size_t>();
                vector_size_t *rawOffsets = offsets->asMutable<vector_size_t>();
                vector_size_t *rawSizes = sizes->asMutable<vector_size_t>();

                vector_size_t newOffset = 0;
                context.applyToSelectedNoThrow(rows, [&](auto row) {
                    vector_size_t arrayRow = decodedArg->index(row);
                    vector_size_t size = array_vector->sizeAt(arrayRow);
                    vector_size_t offset = array_vector->offsetAt(arrayRow);

                    std::iota(rawIndices + newOffset, rawIndices + newOffset + size, offset);
                    std::shuffle(
                        rawIndices + newOffset, rawIndices + newOffset + size, *randGen_);

                    rawSizes[row] = size;
                    rawOffsets[row] = newOffset;
                    newOffset += size;
                });

                auto resultElements = BaseVector::wrap_in_dictionary(
                    nullptr, shuffledIndices, numElements, elementsVector);
                auto localResult = std::make_shared<ArrayVector>(
                    context.pool(),
                    array_vector->type(),
                    nullptr,
                    rows.end(),
                    std::move(offsets),
                    std::move(sizes),
                    std::move(resultElements));

                context.moveOrCopyResult(localResult, rows, result);
            }

        private:
            std::unique_ptr<std::mt19937> randGen_;
        };
    } // namespace

    exec::VectorFunctionMetadata getMetadataForArrayShuffle() {
        return exec::VectorFunctionMetadataBuilder().deterministic(false).build();
    }

    std::vector<std::shared_ptr<exec::FunctionSignature> > arrayShuffleSignatures() {
        return {
            // array(T) -> array(T)
            exec::FunctionSignatureBuilder()
            .typeVariable("T")
            .returnType("array(T)")
            .argumentType("array(T)")
            .build()
        };
    }

    std::shared_ptr<exec::VectorFunction> makeArrayShuffle(
        const std::string &name,
        const std::vector<exec::VectorFunctionArg> &inputArgs,
        const core::QueryConfig &config) {
        return std::make_unique<ArrayShuffleFunction>(std::random_device{}());
    }

    std::vector<std::shared_ptr<exec::FunctionSignature> >
    arrayShuffleWithCustomSeedSignatures() {
        return {
            exec::FunctionSignatureBuilder()
            .typeVariable("T")
            .returnType("array(T)")
            .argumentType("array(T)")
            .constantArgumentType("bigint")
            .build()
        };
    }

    std::shared_ptr<exec::VectorFunction> makeArrayShuffleWithCustomSeed(
        const std::string &name,
        const std::vector<exec::VectorFunctionArg> &inputArgs,
        const core::QueryConfig &config) {
        POLLUX_USER_CHECK_EQ(inputArgs.size(), 2);
        POLLUX_USER_CHECK_EQ(inputArgs[1].type->kind(), TypeKind::BIGINT);
        POLLUX_USER_CHECK_NOT_NULL(inputArgs[1].constantValue);
        POLLUX_CHECK(!inputArgs[1].constantValue->is_null_at(0));

        const auto seed = inputArgs[1]
                .constantValue->template as<ConstantVector<int64_t> >()
                ->value_at(0);
        return std::make_shared<ArrayShuffleFunction>(
            seed + config.sparkPartitionId());
    }
} // namespace kumo::pollux::functions
